#include <iostream>
#include <thread>
#include <vector>
#include <atomic>
#include <chrono>
#include <string>
#include "mq_connection.hpp"
#include "mq_channel.hpp"
#include "mq_worker.hpp"
#include "../mqcommon/mq_helper.hpp"
#include "../mqcommon/mq_logger.hpp"
#include "../mqcommon/mq_msg.pb.h"

// 测试参数
struct TestConfig 
{
    int thread_count;       // 线程数
    int messages_per_thread; // 每个线程发送的消息数
    std::string exchange_name; // 交换机名称
    std::string routing_key;   // 路由键
    bool durable;           // 是否持久化
};

// 测试结果
struct TestResult 
{
    int total_messages;     // 总消息数
    double total_time;      // 总时间(秒)
    double qps;             // 每秒处理消息数
};

// 工作线程函数
void worker_thread(
    MQ::Connection::ptr conn,
    const TestConfig& config,
    std::atomic<int>& completed_count) 
    {

    // 打开通道
    MQ::Channel::ptr channel = conn->openChannel();
    if (!channel) 
    {
        ELOG("打开通道失败");
        return;
    }

    // 准备消息属性
    MQ::BasicProperties bp;
    bp.set_delivery_mode(config.durable ? MQ::DeliveryMode::DURABLE : MQ::DeliveryMode::UNDURABLE);
    bp.set_routing_key(config.routing_key);

    // 发送消息
    for (int i = 0; i < config.messages_per_thread; ++i) 
    {
        bp.set_id(MQ::UUIDHelper::uuid());
        std::string msg_body = "QPS test message - " + std::to_string(i);
        channel->basicPublish(config.exchange_name, &bp, msg_body);
    }

    // 关闭通道
    conn->closeChannel(channel);

    // 增加完成计数
    completed_count += config.messages_per_thread;
}

// 运行测试
TestResult run_test(const TestConfig& config) 
{
    TestResult result;
    result.total_messages = config.thread_count * config.messages_per_thread;

    // 初始化异步工作线程
    MQ::AsyncWorker::ptr awp = std::make_shared<MQ::AsyncWorker>();

    // 创建连接
    MQ::Connection::ptr conn = std::make_shared<MQ::Connection>("127.0.0.1", 8888, awp);

    // 等待连接建立
    std::this_thread::sleep_for(std::chrono::seconds(1));

    // 声明交换机
    MQ::Channel::ptr setup_channel = conn->openChannel();
    if (setup_channel) 
    {
        google::protobuf::Map<std::string, std::string> args;
        setup_channel->declareExchange(
            config.exchange_name,
            MQ::ExchangeType::DIRECT,
            config.durable,
            false,
            args);
        conn->closeChannel(setup_channel);
    }

    // 准备原子计数器
    std::atomic<int> completed_count(0);

    // 记录开始时间
    auto start_time = std::chrono::high_resolution_clock::now();

    // 创建并启动工作线程
    std::vector<std::thread> threads;
    for (int i = 0; i < config.thread_count; ++i) 
    {
        threads.emplace_back(
            worker_thread,
            conn,
            std::ref(config),
            std::ref(completed_count));
    }

    // 等待所有线程完成
    for (auto& t : threads) 
    {
        if (t.joinable()) 
        {
            t.join();
        }
    }

    // 记录结束时间
    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end_time - start_time;
    result.total_time = elapsed.count();
    result.qps = result.total_messages / result.total_time;

    return result;
}

int main(int argc, char* argv[]) 
{
    // 默认测试配置
    TestConfig config;
    config.thread_count = 10;
    config.messages_per_thread = 1000;
    config.exchange_name = "qps_test_exchange";
    config.routing_key = "qps_test_key";
    config.durable = false;

    // 打印测试配置
    ILOG("=== QPS测试配置 ===");
    ILOG("线程数: %d", config.thread_count);
    ILOG("每线程消息数: %d", config.messages_per_thread);
    ILOG("总消息数: %d", config.thread_count * config.messages_per_thread);
    ILOG("交换机名称: %s", config.exchange_name.c_str());
    ILOG("路由键: %s", config.routing_key.c_str());
    ILOG("持久化: %s", config.durable ? "是" : "否");
    ILOG("====================");

    // 运行测试
    TestResult result = run_test(config);

    // 打印测试结果
    ILOG("=== QPS测试结果 ===");
    ILOG("总消息数: %d", result.total_messages);
    ILOG("总时间: %.3f秒", result.total_time);
    ILOG("QPS: %.2f消息/秒", result.qps);
    ILOG("====================");

    return 0;
}